package org.jeecg.modules.airag.llm.document;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.openai.OpenAiTokenizer;

import java.util.ArrayList;
import java.util.List;


class ValidationUtils {
    public static String ensureNotBlank(String value, String name) {
        if (value == null || value.trim().isEmpty()) {
            throw new IllegalArgumentException(name + " must not be blank");
        }
        return value;
    }
    
    public static <T> T ensureNotNull(T value, String name) {
        if (value == null) {
            throw new IllegalArgumentException(name + " must not be null");
        }
        return value;
    }
}

// 工具类包含创建分割器的方法
public class CustomDocumentSplitters {
    
    private static final int DEFAULT_SEGMENT_SIZE = 1000;
    private static final int DEFAULT_OVERLAP_SIZE = 50;
    
    /**
     * 创建支持@符号分割的DocumentSplitter
     * @param tokenizer 分词器
     * @return DocumentSplitter实例
     */
    public static DocumentSplitter atSignSplitter(OpenAiTokenizer tokenizer) {
        return atSignSplitter(DEFAULT_SEGMENT_SIZE, DEFAULT_OVERLAP_SIZE, tokenizer);
    }
    
    /**
     * 创建支持@符号分割的DocumentSplitter
     * @param segmentSize 分段大小
     * @param overlapSize 重叠大小
     * @param tokenizer 分词器
     * @return DocumentSplitter实例
     */
    public static DocumentSplitter atSignSplitter(int segmentSize, int overlapSize, OpenAiTokenizer tokenizer) {
        // 创建递归分割器作为备用
        DocumentSplitter recursiveSplitter = DocumentSplitters.recursive(segmentSize, overlapSize, tokenizer);
        
        // 返回自定义的DocumentSplitter实现
        return new DocumentSplitter() {
            @Override
            public List<TextSegment> split(Document document) {
                List<TextSegment> segments = new ArrayList<>();
                String content = document.text();
                
                // 按@符号分割内容
                String[] parts = content.split("####");
                
                for (String part : parts) {
                    String trimmedPart = part.trim();
                    if (trimmedPart.isEmpty()) {
                        continue;
                    }
                    
                    // 检查当前部分的token数量
//                    if (tokenizer.tokenCount(trimmedPart) > segmentSize) {
//                        // 如果超过限制，使用递归分割器处理
//                        Document tempDoc = new Document(trimmedPart, document.getMetadata());
//                        segments.addAll(recursiveSplitter.split(tempDoc));
//                    } else {
                        // 否则直接创建TextSegment
                        segments.add(new TextSegment(trimmedPart, document.metadata()));
//                    }
                }
                
                return segments;
            }
        };
    }
}

// 使用示例
//class UsageExample {
//    public static void main(String[] args) {
//        OpenAiTokenizer tokenizer = new OpenAiTokenizer();
//
//        // 创建自定义分割器
//        DocumentSplitter splitter = CustomDocumentSplitters.atSignSplitter(tokenizer);
//
//        // 使用分割器
//        Document document = new Document("第一段内容@第二段内容@这是一段很长的内容...");
//        List<TextSegment> result = splitter.split(document);
//    }
//}
    